from fusionmatrix.base.widgets import ProgressDisplayer
from PySide6.QtCore import QTimer

from multiprocessing import Process, Queue, set_start_method
from pydantic import BaseModel, Field
from queue import Empty
from typing import List, Set, Optional, ClassVar, Dict
from loguru import logger
from time import sleep


set_start_method('spawn', True)


def worker_run(idx: int, task_queue: Queue, result_queue: Queue) -> None:
    from fusionmatrix.analyzers.classifier.base import ClassificationTask

    logger.info(f"worker {idx} started")
    while True:
        try:
            task = task_queue.get_nowait()
        except Empty:
            sleep(0.2)
            continue

        logger.info(f"worker {idx} accept task")
        if isinstance(task, tuple) and task[0] == 'ClassificationTask':
            task = ClassificationTask.model_validate_json(task[1])
        else:
            logger.error(f"unknown task type: {task}")
            continue

        result = task.run()
        result_queue.put(('TaskResult', result.model_dump_json()))


class TaskResult(BaseModel):
    task_id: int
    task_status: bool
    result: str = ''
    message: str = ''


class BaseTask(BaseModel):
    task_id: int = -1
    prev_ids: Set[int] = Field(default_factory=set)
    cnt: ClassVar[int] = 0

    def initializeTaskId(self) -> None:
        BaseTask.cnt += 1
        self.task_id = BaseTask.cnt

    def run(self) -> TaskResult:
        """ The function is expected to run at the worker process """
        raise NotImplementedError

    def notice(self, result: TaskResult) -> None:
        """ The function is expected to run at the main process """
        raise NotImplementedError


class TaskWorker:

    def __init__(self, idx: int, task_queue: Queue, result_queue: Queue) -> None:
        self.task_queue = task_queue
        self.result_queue = result_queue
        self.idx = idx
        self.proc: Optional[Process] = None

    def start(self) -> None:
        if self.proc is not None:
            return

        proc = Process(target=worker_run, args=(self.idx, self.task_queue, self.result_queue))
        self.proc = proc
        self.proc.start()

    def stop(self) -> None:
        if self.proc is None:
            return

        logger.info(f"kill worker {self.idx}")
        self.proc.terminate()
        self.proc.join()
        self.proc = None


class TaskExecutor:
    task_queue: Queue = Queue()
    result_queue: Queue = Queue()
    workers: List[TaskWorker] = []
    tasks: Dict[int, BaseTask] = {}
    nfinished: int = 0
    progress_displayer: Optional[ProgressDisplayer] = None
    timer: QTimer = QTimer(None)

    @classmethod
    def initialize(cls) -> None:
        cls.timer.setInterval(100)
        cls.timer.setSingleShot(False)
        cls.timer.timeout.connect(cls.checkResult)

    @classmethod
    def shutdown(cls) -> None:
        cls.clearExecutor()

    @classmethod
    def createWorker(cls) -> None:
        worker = TaskWorker(len(cls.workers) + 1, cls.task_queue, cls.result_queue)
        cls.workers.append(worker)

    @classmethod
    def setWorkerNumber(cls, worker_number: int) -> None:
        while len(cls.workers) < worker_number:
            cls.createWorker()

    @classmethod
    def submitTask(cls, task: BaseTask) -> None:
        cls.submitBatchTasks([task])

    @classmethod
    def submitBatchTasks(cls, tasks: List[BaseTask]) -> None:
        cls.startAllWorkers()
        for task in tasks:
            task.initializeTaskId()
            cls.tasks[task.task_id] = task
            cls.task_queue.put((task.__class__.__name__, task.model_dump_json()))

    @classmethod
    def startAllWorkers(cls) -> None:
        cls.timer.start()
        for worker in cls.workers:
            worker.start()

    @classmethod
    def stopAllWorkers(cls) -> None:
        cls.timer.stop()
        for worker in cls.workers:
            worker.stop()

        cls.nfinished = 0
        cls.tasks.clear()
        if cls.progress_displayer is not None:
            cls.progress_displayer.hide()

    @classmethod
    def clearExecutor(cls) -> None:
        cls.stopAllWorkers()
        while not cls.task_queue.empty():
            _ = cls.task_queue.get_nowait()
        while not cls.result_queue.empty():
            _ = cls.result_queue.get_nowait()

    @classmethod
    def checkResult(cls) -> None:
        while True:
            try:
                classname, payload = cls.result_queue.get_nowait()
            except Empty:
                break

            if classname == 'TaskResult':
                result = TaskResult.model_validate_json(payload)
                task = cls.tasks[result.task_id]
                task.notice(result)
                cls.nfinished += 1

        if cls.progress_displayer is not None:
            cls.progress_displayer.setProgress(cls.nfinished, len(cls.tasks))
            if cls.nfinished == len(cls.tasks):
                cls.stopAllWorkers()
